// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#define PLATFORM_CUDA

// clang-format off
#define CUDA_PLATFORM_SPECIALIZATION_HEADER \
  breeze/platforms/specialization/cuda-ptx.cuh
// clang-format on

#include <breeze/functions/reduce.h>
#include <breeze/platforms/platform.h>
#include <breeze/utils/device_vector.h>
#include <breeze/platforms/cuda.cuh>

#include <gtest/gtest.h>

namespace breeze {
namespace {

using namespace functions;
using namespace utils;

constexpr int kBlockThreads = 256;
constexpr int kItemsPerThread = 8;
constexpr int kBlockItems = kBlockThreads * kItemsPerThread;
constexpr int kNumItems = 250'000;
constexpr int kNumBlocks = (kNumItems + kBlockItems - 1) / kBlockItems;

__global__ __launch_bounds__(kBlockThreads) void reduceKernel(int* out) {
  CudaPlatform<kBlockThreads, 32> p;
  using BlockReduceT = BlockReduce<decltype(p), int>;
  __shared__ typename BlockReduceT::Scratch scratch;
  int items[kItemsPerThread];
  for (int i = 0; i < kItemsPerThread; ++i) {
    items[i] = 1;
  }
  int aggregate = BlockReduceT::template Reduce<ReduceOpAdd, kItemsPerThread>(
      p,
      make_slice(items),
      make_slice(&scratch).template reinterpret<SHARED>());
  if (p.thread_idx() == 0) {
    out[p.block_idx()] = aggregate;
  }
}

TEST(BreezeCudaTest, reduce) {
  device_vector<int> result(kNumBlocks);
  reduceKernel<<<kNumBlocks, kBlockThreads>>>(result.data());
  std::vector<int> actual(kNumBlocks);
  result.copy_to_host(actual.data(), actual.size());
  std::vector<int> expected(kNumBlocks, kBlockThreads * kItemsPerThread);
  ASSERT_EQ(actual, expected);
}

} // namespace
} // namespace breeze
